package com.LeeCode;

import java.util.*;

/**
 * 设计路由器
 */

public class Code3508 {
    public static void main(String[] args) {
        Router router = new Code3508().new Router(3);
        System.out.println(router.addPacket(1, 4, 90));
        System.out.println(router.addPacket(2, 5, 90));
        System.out.println(router.addPacket(1, 4, 90));
        System.out.println(router.addPacket(3, 5, 95));
        System.out.println(router.addPacket(4, 5, 105));
        router.forwardPacket();
        System.out.println(router.addPacket(5, 2, 110));
        System.out.println(router.getCount(5, 100, 110));

    }

    class Router {
        class Packet {
            public int source;
            public int destination;
            int timestamp;

            public Packet(int source, int destination, int timestamp) {
                this.source = source;
                this.timestamp = timestamp;
                this.destination = destination;
            }

            @Override
            public boolean equals(Object o) {
                if (this == o) return true;
                if (!(o instanceof Packet)) return false;

                Packet other = (Packet) o;

                return source == other.source &&
                        destination == other.destination &&
                        timestamp == other.timestamp;
            }

            @Override
            public int hashCode() {
                // 使用 Objects.hash 可以简化 hash 的计算
                return Objects.hash(source, destination, timestamp);
            }
        }

        class Pair {
            List<Integer> timestamps;
            int head;

            public Pair(List<Integer> timestamps, int head) {
                this.timestamps = timestamps;
                this.head = head;
            }
        }

        private int memoryLimit;
        private final Queue<Packet> packetQ = new ArrayDeque<>();
        private final Set<Packet> packetSet = new HashSet<>();
        private final Map<Integer, Pair> destmap = new HashMap<>();

        public Router(int memoryLimit) {
            this.memoryLimit = memoryLimit;
        }

        public boolean addPacket(int source, int destination, int timestamp) {
            Packet packet = new Packet(source, destination, timestamp);
            if (!packetSet.add(packet)) {
                return false;
            }
            if (packetQ.size() == memoryLimit) {
                forwardPacket();
            }

            packetQ.add(packet);
            destmap.computeIfAbsent(destination, k -> new Pair(new ArrayList<>(), 0)).timestamps.add(timestamp);
            return true;
        }

        public int[] forwardPacket() {
            if (packetQ.isEmpty()) {
                return new int[]{};
            }
            Packet packet = packetQ.poll();
            packetSet.remove(packet);
            Pair pair = destmap.get(packet.destination);
            pair.head += 1;
            return new int[]{packet.source, packet.destination, packet.timestamp};
        }

        public int getCount(int destination, int startTime, int endTime) {
            Pair pair = destmap.get(destination);
            if (pair == null) {
                return 0;
            }
            int left = lowerBound(pair.timestamps, startTime, pair.head - 1);
            int right = lowerBound(pair.timestamps, endTime + 1, pair.head - 1);
            return right - left;
        }

        private int lowerBound(List<Integer> nums, int target, int left) {
            int right = nums.size();
            while (left + 1 < right) {
                int mid = left + (right - left) / 2;
                if (nums.get(mid) >= target) {
                    right = mid;
                } else {
                    left = mid;
                }
            }
            return right;
        }
    }

}
